import numpy as np
import pandas as pd
from tqdm import *
from sklearn.metrics import mutual_info_score
from scipy.stats import entropy
import matplotlib.pyplot as plt
def Mutal_information(trace_url,x_url,blocknum,types,samples,pre_url):
    mi = np.zeros(samples)
    x_all = np.zeros(5000*blocknum)
    trace_all = np.zeros([5000*blocknum,samples])
    for block in trange(blocknum,desc="Read files"):
        x = np.load(x_url.format(block))&0xff
        trace = np.load(trace_url.format(block))
        for i in range(5000):
            x_all[5000*block+i] = x[i]
            trace_all[5000 * block + i] = trace[i]
    X = x_all.tolist()
    for num in trange(samples,desc="Computing Mutal information"):
        Y = trace_all[:,num]
        num_bins = 10
        Y = np.digitize(Y, bins=np.linspace(min(Y), max(Y), num_bins))
        # Y = pd.qcut(Y, q=num_bins, labels=False, duplicates='drop')
        mi[num]=mutual_info_score(X,Y)
    np.save(pre_url + r"\mi-result.npy", mi)
    plt.rcParams['figure.figsize'] = (12.0, 8.0)
    plt.plot(mi)
    plt.rcParams['figure.figsize'] = (12.0, 8.0)
    plt.xlabel('Trace sample')
    plt.ylabel('Mutual Information')
    plt.title('MI result')
    plt.show()

if __name__ =="__main__":
    pre_url = r"E:\CBM\traces\snr-base-core-cbm-sbox"
    trace= pre_url + r"\trace-batch\tracebatch-{}.npy"
    x= pre_url + r"\x0-batch\x0batch-{}.npy"
    Mutal_information(trace,x,60,types=256,samples=10000,pre_url=pre_url)
    